Hierarchical Bayesian Neural Networks with variable group sizes

Warning

This entry is seriously in the making. Please come back later for updates.

Bayesian Modeling is the primary choice if you want to obtain the uncertainty associated with the predictions of a model. Of course, there are also voices arguing why Bayesian Deep Learning is a promising avenue.

There are some very useful blog entries and notebooks out there (e.g. by Thomas Wiecki using Theano PyMC3 and this repo using a more recent version of JAX).

However, those examples only work with the critical assumption that the group sizes are all of the same size. In reality, this is rarely the case, of course.

Here, I will show you how you can implement a Hierarchical Bayesian Neural Network irrespective of the group sizes you observe in your dataset.

from typing import Tuple
from datetime import date
from functools import partial
from warnings import filterwarnings
import jax
import jax.random as jr
import jax.numpy as jnp
import equinox as eqx
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons
from sklearn.preprocessing import scale
import blackjax
import tensorflow_probability.substrates.jax.distributions as tfd
from sklearn.model_selection import train_test_split


filterwarnings("ignore")

key = jr.key(int(date.today().strftime("%Y%m%d")))
cmap = mpl.colormaps["PiYG"]

WARNING:2025-06-11 22:20:45,725:jax._src.xla_bridge:791: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

# Data
n_groups = 18
n_grps_sq = int(np.sqrt(n_groups))
n_samples = np.random.randint(10, 200, size=n_groups)
noise = 0.3

# MLP params
data_dim = 2
hidden_layer_width = 8
n_hidden_layers = 3

# Sampling
num_warmup = 200
num_samples = 400

X, Y = make_moons(noise=noise, n_samples=1000)

for i in range(2):
    plt.scatter(X[Y == i, 0], X[Y == i, 1], color=cmap(float(i)), label=f"Class {i}", alpha=0.8)
plt.legend()
plt.show()

The standard two-moons dataset plot

def rotate(X, deg):
    theta = np.radians(deg)
    c, s = np.cos(theta), np.sin(theta)
    R = np.matrix([[c, -s], [s, c]])
    X = X.dot(R)
    return np.asarray(X)


np.random.seed(31)

Xs, Ys, gs = [], [], []
Xs_train, Ys_train, gs_train, Xs_test, Ys_test, gs_test = [], [], [], [], [], []

for i in range(n_groups):
    # Generate data with 2 classes that are not linearly separable
    X, Y = make_moons(noise=noise, n_samples=n_samples[i])
    X = scale(X)

    # Rotate the points randomly for each category
    rotate_by = np.random.randn() * 90.0
    X = rotate(X, rotate_by)
    Xs.append(X)
    Ys.append(Y)
    gs.append(X.shape[0])
    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=31)
    Xs_train.append(X_train)
    Ys_train.append(Y_train)
    gs_train.append(X_train.shape[0])
    Xs_test.append(X_test)
    Ys_test.append(Y_test)
    gs_test.append(X_test.shape[0])

def pad_arrays(arrays, fill_value):
    max_size = max(array.shape[0] for array in arrays)
    padded_arrays = []
    for array in arrays:
        if array.ndim == 1:
            padding = (0, max_size - array.shape[0]) 
            padded_array = jnp.pad(array, padding, mode="constant", constant_values=fill_value)
            padded_arrays.append(padded_array[:, np.newaxis])  
        else:
            padding = ((0, max_size - array.shape[0]), (0, 0))  
            padded_array = jnp.pad(array, padding, mode="constant", constant_values=fill_value)
            padded_arrays.append(padded_array)
    return padded_arrays

# Stack group arrays and create a mask
fill_value = 1e5
Xs_train = jnp.stack(pad_arrays(Xs_train, fill_value))
Ys_train = jnp.stack(pad_arrays(Ys_train, fill_value)).squeeze(axis=2)
Xs_test = jnp.stack(pad_arrays(Xs_test, fill_value))
Ys_test = jnp.stack(pad_arrays(Ys_test, fill_value)).squeeze(axis=2)

mask_train = jnp.where(Xs_train == fill_value, fill_value, 1)
mask_test = jnp.where(Xs_test == fill_value, fill_value, 1)

# Number of rows and columns for subplots
n_rows = int(np.ceil(n_groups / n_grps_sq))
n_cols = n_grps_sq

fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 10), sharex=True, sharey=True)

# Flatten axes array for easy iteration
axes = axes.flatten()

for i, (X, Y, group_size, ax) in enumerate(zip(Xs_train, Ys_train, gs_train, axes)):
    for c in range(2):
        ax.scatter(X[Y == c, 0], X[Y == c, 1], color=cmap(float(c)), label=f"Class {c}", alpha=0.8)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.legend(frameon=False)
    ax.set(title=f"Category {i + 1}, N_training = {group_size}")

# Hide any unused subplots
for j in range(n_groups, len(axes)):
    fig.delaxes(axes[j])

plt.tight_layout()
plt.show()

18 groups with varying numbers of training examples

grid = jnp.mgrid[-3:3:100j, -3:3:100j].reshape((2, -1)).T
grid_3d = jnp.repeat(grid[None, ...], n_groups, axis=0)
mask_grid = jnp.ones(grid_3d.shape)

@eqx.filter_vmap
def make_ensemble(key):
    # Create an ensemble of models
    net = eqx.nn.MLP(in_size=data_dim, out_size=1, width_size=hidden_layer_width, depth=n_hidden_layers, key=key)
    return net


@eqx.filter_vmap(in_axes=(eqx.if_array(0), None))
def evaluate_ensemble(ensemble, x):
    # Evaluate each member of the ensemble on the same data
    o = ensemble(x)
    return o.mean()


def evaluate_per_ensemble(model, x):
    return jax.vmap(model)(x)


def apply_ensemble(ensemble, D):
    # ensemble_fn = partial(evaluate_ensemble, ensemble)
    preds = eqx.filter_vmap(evaluate_per_ensemble)(ensemble, D)
    return preds


key = jr.PRNGKey(0)
hnn = make_ensemble(jr.split(key, n_groups))

class NonCentredLinear(eqx.Module):
    # mu is the global network
    mu: jax.Array
    # eps are the n_group local networks
    eps: jax.Array
    std: jax.Array

    def __init__(self, in_size, out_size, n_groups, *, key):
        self.mu = jr.normal(key, (in_size, out_size))
        self.eps = jr.normal(key, (n_groups, in_size, out_size))
        self.std = jnp.ones((1,))

    def __call__(self, x):
        w = self.mu + self.std * self.eps
        return x @ w


class HNN(eqx.Module):
    layers: Tuple[NonCentredLinear]
    out: eqx.nn.Linear

    def __init__(self, layer_width, n_layers, n_groups, *, key):
        dims = [data_dim] + [layer_width] * n_layers
        layers = []
        for n, (_in, _out) in enumerate(zip(dims[:-1], dims[1:])):
            layer = NonCentredLinear(_in, _out, n_groups, key=jr.fold_in(key, n))
            layers += [layer]
        self.layers = tuple(layers)
        self.out = eqx.nn.Linear(layer_width, 1, key=key)

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
            x = jax.nn.tanh(x)
        # Vmap over groups and samples
        o = jax.vmap(jax.vmap(self.out))(x)
        return o

def get_init_apply_fns(model):
    params, static = eqx.partition(model, eqx.is_inexact_array)

    def init_fn():
        return params

    def apply_fn(_params, x):
        model = eqx.combine(_params, static)
        return model(x)

    return init_fn, apply_fn

hnn = HNN(hidden_layer_width, n_hidden_layers, n_groups, key=key)
init_fn, apply_fn = get_init_apply_fns(hnn)
params = init_fn()

def inference_loop(key, step_fn, initial_state, num_samples):
    def one_step(state, key):
        state, _ = step_fn(key, state)
        return state, state

    keys = jr.split(key, num_samples)
    _, states = jax.lax.scan(one_step, initial_state, keys)
    return states

def get_predictions(model, samples, X, mask, fill_value, key):
    vectorized_apply = jax.vmap(apply_fn, in_axes=(0, None), out_axes=0)
    z = vectorized_apply(samples, X)
    predictions = tfd.Bernoulli(logits=z).sample(seed=key)
    mask_reshaped = jnp.broadcast_to(
        jnp.mean(mask, axis=-1).reshape(mask.shape[0], mask.shape[1], 1), predictions.shape
    )
    predictions = jnp.where(mask_reshaped == fill_value, jnp.nan, predictions)
    return predictions.squeeze(-1)

def get_mean_predictions(predictions, threshold=0.5):
    # Compute mean prediction and confidence interval around median
    mean_prediction = jnp.nanmean(predictions, axis=0)
    return mean_prediction > threshold

def fit_and_eval(
    key,
    initial_position,  # Passed from `init_fn` of init/apply function conversion of Equinox NN
    model,
    logdensity_fn,
    X_train,
    Y_train,
    mask_train,
    fill_value,
    X_test,
    grid,
    num_warmup=200,
    num_samples=100,
):
    (
        warmup_key,
        inference_key,
        train_key,
        test_key,
        grid_key,
    ) = jr.split(key, 5)

    # Initialization
    logprob = partial(logdensity_fn, X=X_train, Y=Y_train, mask=mask_train, fill_value=fill_value, model=model)

    # Warm up
    adapt = blackjax.window_adaptation(blackjax.nuts, logprob)
    (final_state, params), _ = adapt.run(warmup_key, initial_position, num_warmup)
    step_fn = blackjax.nuts(logprob, **params).step

    # Inference
    states = inference_loop(inference_key, step_fn, final_state, num_samples)
    samples = states.position

    # Evaluation
    predictions = get_predictions(model, samples, X_train, mask_train, fill_value, train_key)
    Y_pred_train = get_mean_predictions(predictions)

    predictions = get_predictions(model, samples, X_test, mask_test, fill_value, test_key)
    Y_pred_test = get_mean_predictions(predictions)

    pred_grid = get_predictions(model, samples, grid, mask_grid, fill_value, grid_key)

    return Y_pred_train, Y_pred_test, pred_grid

def logprior_fn(params):
    normal = tfd.Normal(0.0, 1.0)
    leaves, _ = jax.tree_util.tree_flatten(params)
    flat_params = jnp.concatenate([jnp.ravel(a) for a in leaves])
    return jnp.sum(normal.log_prob(flat_params))

def logprior_fn_of_hnn(params, model):
    """p(w) where w is NN(X; w)"""
    lp = 0.0
    half_normal = tfd.HalfNormal(1.0)
    normal = tfd.Normal(0.0, 1.0)
    for layer in params.layers:
        lp += normal.log_prob(layer.mu).sum()
        lp += normal.log_prob(layer.eps).sum()
        lp += half_normal.log_prob(layer.std).sum()
    lp += logprior_fn(params.out)
    return lp


def loglikelihood_fn(params, X, Y, mask, fill_value, model):
    """p(Y|Y_=NN(X; w))"""
    logits = jnp.ravel(apply_fn(params, X))
    # apply the mask: where the mask has the fill value, the logits should also be zero
    logits = jnp.where(jnp.ravel(mask[:, :, 0]) == fill_value, 0, logits)
    return jnp.sum(tfd.Bernoulli(logits).log_prob(jnp.ravel(Y)))


def logdensity_fn_of_hnn(params, X, Y, mask, fill_value, model):
    return logprior_fn_of_hnn(params, model) + loglikelihood_fn(params, X, Y, mask, fill_value, model)

key, inference_key = jr.split(key)

(Ys_hierarchical_pred_train, Ys_hierarchical_pred_test, ppc_grid) = fit_and_eval(
    inference_key,
    params,
    hnn,
    logdensity_fn_of_hnn,
    Xs_train,
    Ys_train,
    mask_train,
    fill_value,
    Xs_test,
    grid_3d,
    num_warmup=num_warmup,
    num_samples=num_samples,
)

def reverse_mask(targets, predictions, mask, fill_value):
    targets, predictions, mask = jnp.ravel(targets), jnp.ravel(predictions), jnp.ravel(mask[:,:,0])
    positions_to_omit = jnp.where(mask == fill_value)[0]
    filtered_targets, filtered_predictions = jnp.delete(targets, positions_to_omit), jnp.delete(predictions, positions_to_omit)
    return filtered_targets,filtered_predictions    

filtered_Ys_train, filtered_Ys_hierarchical_pred_train = reverse_mask(
    Ys_train, Ys_hierarchical_pred_train, mask_train, fill_value
)
filtered_Ys_test, filtered_Ys_hierarchical_pred_test = reverse_mask(
    Ys_test, Ys_hierarchical_pred_test, mask_test, fill_value
)

print("Train accuracy = {:.2f}%".format(100 * jnp.mean(filtered_Ys_hierarchical_pred_train == filtered_Ys_train)))

Train accuracy = 92.48%

print("Test accuracy = {:.2f}%".format(100 * jnp.mean(filtered_Ys_hierarchical_pred_test == filtered_Ys_test)))

Test accuracy = 89.02%

def plot_decision_surfaces_hierarchical(nrows=2, ncols=2):
    fig, axes = plt.subplots(figsize=(15, 12), nrows=nrows, ncols=ncols, sharex=True, sharey=True)

    for i, (X, Y_pred, Y_true, ax) in enumerate(zip(Xs_train, Ys_hierarchical_pred_train, Ys_train, axes.flatten())):
        ax.contourf(
            grid[:, 0].reshape((100, 100)),
            grid[:, 1].reshape((100, 100)),
            ppc_grid[:, i, :].mean(axis=0).reshape(100, 100),
            cmap=cmap,
            zorder=0,
        )
        for i in range(2):
            ax.scatter(X[Y_true == i, 0], X[Y_true == i, 1], color="w", alpha=0.8, s=20.0, zorder=1)
            ax.scatter(
                X[Y_true == i, 0],
                X[Y_true == i, 1],
                color=cmap(float(i)),
                label=f"Class {i}",
                alpha=0.8,
                s=10.0,
                zorder=2,
            )
        ax.set_xticks([])
        ax.set_yticks([])
        ax.legend(frameon=False)


# %%
plot_decision_surfaces_hierarchical(nrows=n_grps_sq, ncols=n_grps_sq)
plt.savefig("figures/hbnn_decision_boundaries.png", bbox_inches="tight")
plt.show()